// Copyright 2022 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef THIRD_PARTY_DNN_ACCELERATOR_HLS_SYSTEMC_SKEWER_H_
#define THIRD_PARTY_DNN_ACCELERATOR_HLS_SYSTEMC_SKEWER_H_

#include <mc_connections.h>
#include <systemc.h>

#include <boost/preprocessor/stringize.hpp>
#include <boost/preprocessor/repetition/repeat.hpp>

#include "src/AccelTypes.h"
#include "src/ArchitectureParams.h"

#define REPEAT(x) BOOST_PP_REPEAT(DIMENSION, x, 0)

template <typename T, int NUM_REGS>
class Fifo {
 public:
  Fifo() {}

  void run(T &input, T &output) {
    for (int i = NUM_REGS - 1; i >= 0; i--) {
      if (i == 0) {
        regs[i] = input;
      } else {
        regs[i] = regs[i - 1];
      }

      output = regs[NUM_REGS - 1];
    }
  }

 private:
  T regs[NUM_REGS];
};

/*
 * Takes an input of Pack1D<DTYPE,SIZE> and skews it to produce n=SIZE outputs
 * of DTYPE
 */
template <typename DTYPE, int SIZE>
SC_MODULE(SerializedSkewer) {
 private:
#define DECL_FIFOS(z, i, unused) sc_fifo<DTYPE> BOOST_PP_CAT(fifo, i);
  REPEAT(DECL_FIFOS)
#undef DECL_FIFOS
  int dummy;

 public:
  sc_in<bool> CCS_INIT_S1(clk);
  sc_in<bool> CCS_INIT_S1(rstn);

  Connections::In<Pack1D<DTYPE, SIZE> > CCS_INIT_S1(din);
  Connections::Out<DTYPE> dout[SIZE];

#define FIFO_SIZE_INIT(z, i, unused) BOOST_PP_CAT(fifo, i)(i + 2),

  SC_CTOR(SerializedSkewer) : REPEAT(FIFO_SIZE_INIT) dummy(0) {
    SC_THREAD(writeFifos);
    sensitive << clk.pos();
    async_reset_signal_is(rstn, false);

    // declare threads for reading from fifos
#define DECL_THREADS(z, i, unused)                                          \
  declare_thread_process(BOOST_PP_CAT(BOOST_PP_CAT(readFifos, i), _handle), \
                         BOOST_PP_STRINGIZE(BOOST_PP_CAT(readFifos, i)),    \
                                            SC_CURRENT_USER_MODULE,         \
                                            BOOST_PP_CAT(readFifos, i));    \
  sensitive << clk.pos();                                                   \
  async_reset_signal_is(rstn, false);

    REPEAT(DECL_THREADS)
#undef DECL_THREADS
  }

#undef FIFO_SIZE_INIT

  void writeFifos() {
    din.Reset();
    wait();

#pragma hls_pipeline_init_interval 1
#pragma hls_pipeline_stall_mode flush
    while (true) {
      Pack1D<DTYPE, SIZE> input = din.Pop();

#define FIFO_WRITE(z, i, unused) BOOST_PP_CAT(fifo, i).write(input[i]);
      REPEAT(FIFO_WRITE)
#undef FIFO_WRITE
    }
  }

#define DECL_FUNCS(z, i, unused)                                \
  void BOOST_PP_CAT(readFifos, i)() {                           \
    dout[i].Reset();                                            \
    wait();                                                     \
    _Pragma("hls_pipeline_init_interval 1")                     \
        _Pragma("hls_pipeline_stall_mode flush") while (true) { \
      dout[i].Push(BOOST_PP_CAT(fifo, i).read());               \
    }                                                           \
  }

  REPEAT(DECL_FUNCS)
#undef DECL_FUNCS
};

template <typename DTYPE, int SIZE>
SC_MODULE(DeserializedSkewer) {
 private:
#define DECL_FIFOS(z, i, unused) sc_fifo<DTYPE> BOOST_PP_CAT(fifo, i);
  REPEAT(DECL_FIFOS)
#undef DECL_FIFOS
  int dummy;

 public:
  sc_in<bool> CCS_INIT_S1(clk);
  sc_in<bool> CCS_INIT_S1(rstn);

  Connections::In<DTYPE> din[SIZE];
  Connections::Out<Pack1D<DTYPE, SIZE> > CCS_INIT_S1(dout);

#define FIFO_SIZE_INIT(z, i, unused) BOOST_PP_CAT(fifo, i)(DIMENSION - i + 1),

  SC_CTOR(DeserializedSkewer) : REPEAT(FIFO_SIZE_INIT) dummy(0) {
    SC_THREAD(readFifos);
    sensitive << clk.pos();
    async_reset_signal_is(rstn, false);

    // declare threads for writing to fifos
#define DECL_THREADS(z, i, unused)                                           \
  declare_thread_process(BOOST_PP_CAT(BOOST_PP_CAT(writeFifos, i), _handle), \
                         BOOST_PP_STRINGIZE(BOOST_PP_CAT(writeFifos, i)),    \
                                            SC_CURRENT_USER_MODULE,          \
                                            BOOST_PP_CAT(writeFifos, i));    \
  sensitive << clk.pos();                                                    \
  async_reset_signal_is(rstn, false);

    REPEAT(DECL_THREADS)
#undef DECL_THREADS
  }
#undef FIFO_SIZE_INIT

  void readFifos() {
    dout.Reset();
    wait();

#pragma hls_pipeline_init_interval 1
#pragma hls_pipeline_stall_mode flush
    while (true) {
      Pack1D<DTYPE, SIZE> output;

#define FIFO_READ(z, i, unused) output[i] = BOOST_PP_CAT(fifo, i).read();
      REPEAT(FIFO_READ)
#undef FIFO_READ

      dout.Push(output);
    }
  }

#define DECL_FUNCS(z, i, unused)                                \
  void BOOST_PP_CAT(writeFifos, i)() {                          \
    din[i].Reset();                                             \
    wait();                                                     \
    _Pragma("hls_pipeline_init_interval 1")                     \
        _Pragma("hls_pipeline_stall_mode flush") while (true) { \
      BOOST_PP_CAT(fifo, i).write(din[i].Pop());                \
    }                                                           \
  }
  REPEAT(DECL_FUNCS)
#undef DECL_FUNCS
};

#endif  // THIRD_PARTY_DNN_ACCELERATOR_HLS_SYSTEMC_SKEWER_H_
